#include <iostream>
#include <string>
#include <vector>
using namespace std;

#include "conv.h"
#include "utils.h"
#include "timer.h"
#include "gemm.h"

Conv::Conv(vector<float>& weight, int ic, int oc, int ks, int stride, int padding){
    this->weight = weight;
    this->ks = ks;
    this->stride = stride;
    this->padding = padding;
    this->ic = ic;
    this->oc = oc;
    this->timer_conv.m_name = "conv";
    this->timer_pad.m_name = "pad";
    this->timer_im2col.m_name = "im2col";
    this->timer_gemm.m_name = "gemm";
}

void Conv::forward_direct(vector<float>& input, vector<int> input_shape, vector<float>& output){
    vector<float> padded;
    timer_pad.start();
    pad(input, input_shape, padding, padding, padding, padding, padded);
    timer_pad.pause();

    int ic = input_shape[0];
    int ih = input_shape[1] + 2*padding;
    int iw = input_shape[2] + 2*padding;
    int oh = (ih - ks + 1) / stride;
    int ow = (iw - ks + 1) / stride;

    size_t out_size = oh*ow*oc;
    output.resize(out_size);
    fill(output.begin(), output.end(), 0);
    
    timer_conv.start();
    for(int oc_i=0; oc_i<oc; oc_i++){
        for(int ic_i=0; ic_i<ic; ic_i++){
            for(int outh_i=0; outh_i<oh; outh_i+=stride){
                for(int outw_i=0; outw_i<ow; outw_i+=stride){
                    for(int kh_i=0; kh_i<ks; kh_i++){
                        for(int kw_i=0; kw_i<ks; kw_i++){
                            size_t out_idx = oc_i * oh * ow + outh_i * ow + outw_i;
                            size_t kernel_idx = oc_i *ic*ks*ks + ic_i * ks*ks + kh_i * ks + kw_i;
                            int ih_i = outh_i + kh_i;
                            int iw_i = outw_i + kw_i;
                            size_t inp_idx = ic_i * ih * iw + ih_i * iw + iw_i;
                            output[out_idx] += weight[kernel_idx] * padded[inp_idx];
                        }
                    }
                }
            }
        }
    }
    timer_conv.pause();
}

/*
add some matrix index cache
*/
void Conv::forward_direct1(vector<float>& input, vector<int> input_shape, vector<float>& output){

    vector<float> padded;
    timer_pad.start();
    pad(input, input_shape, padding, padding, padding, padding, padded);
    timer_pad.pause();

    int ic = input_shape[0];
    int ih = input_shape[1] + 2*padding;
    int iw = input_shape[2] + 2*padding;
    int oh = (ih - ks) / stride + 1;
    int ow = (iw - ks) / stride + 1;

    size_t out_size = oh*ow*oc;
    output.resize(out_size);
    fill(output.begin(), output.end(), 0);
    
    timer_conv.start();
    size_t ksks = ks * ks;
    for(int oc_i=0; oc_i<oc; oc_i++){
        size_t out_1 = oc_i * oh * ow;
        size_t k_1 = oc_i *ic*ksks;
        for(int ic_i=0; ic_i<ic; ic_i++){
            size_t inp_1 = ic_i * ih * iw;
            size_t k_2 = ic_i * ksks;
            for(int outh_i=0; outh_i<oh; outh_i+=stride){
                size_t out_2 = out_1 + outh_i * ow;
                for(int outw_i=0; outw_i<ow; outw_i+=stride){
                    size_t out_idx = out_2 + outw_i;
                    for(int kh_i=0; kh_i<ks; kh_i++){
                        size_t k_3 = k_1 + k_2 + kh_i * ks;
                        size_t inp_2 = inp_1 + (outh_i + kh_i) * iw;
                        for(int kw_i=0; kw_i<ks; kw_i++){
                            size_t kernel_idx = k_3 + kw_i;
                            size_t inp_idx = inp_2 + outw_i + kw_i;
                            output[out_idx] += weight[kernel_idx] * padded[inp_idx];
                        }
                    }
                }
            }
        }
    }
    timer_conv.pause();
}


void Conv::forward_im2col_gemm(vector<float>& input, vector<int> input_shape, vector<float>& output){
    int ic = input_shape[0];
    int ih = input_shape[1];
    int iw = input_shape[2];
    
    timer_im2col.start();
    vector<float> B;
    im2col(input, ic, ih, iw, this->ks, this->stride, this->padding, B);
    timer_im2col.pause();
    // B shape: [ic*ks*ks, oh, ow]

    vector<float>& A = this->weight; // [oc, ic*ks*ks]
    
    int oc = this->oc;
    int oh = (ih + 2 * padding - ks) / stride + 1;
    int ow = (iw + 2 * padding - ks) / stride + 1;
    output.resize(oc * oh * ow);
    fill(output.begin(), output.end(), 0);
    
    int M = oc;
    int K = ic * ks * ks;
    int N = oh * ow;

    timer_gemm.start();
    gemm(M, N, K, A.data(), B.data(), output.data());
    // gemm(M, N, K, A, B, output); 
    timer_gemm.pause();
    // output: [oc, oh, ow]
}